import torch.cuda
from torch.backends import cudnn

import config
from util.utils import get_classes
from data import load
from tqdm import tqdm
from util.output import *
from model.model import *
from model.frcnn_loss import *


def _train(model, train_data, valid_data):
    model_train = model.train()
    model_train = torch.nn.DataParallel(model_train)
    cudnn.benchmark = True
    model_train = model_train.to()

    torch.cuda.empty_cache()
    optimizer = config.optimizer(model.parameters(), config.lr, weight_decay=config.weight_decay)
    # 使用Cos变化的学习率
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=1, T_mult=2, eta_min=5e-5)

    # 定义损失函数
    losser = FRCNNLoss(model_train, optimizer)

    loss_list, acc_list = [], []
    bar = tqdm(range(config.epochs))
    for _ in bar:
        model.train()
        total_loss = 0
        rpn_loc_loss = 0
        rpn_cls_loss = 0
        roi_loc_loss = 0
        roi_cls_loss = 0
        for batch in train_data:
            rpn_loc, rpn_cls, roi_loc, roi_cls, total = model.train_step(losser, batch)
            total.backward()

            nn.utils.clip_grad_value_(model.parameters(), config.grad_clip)

            total_loss += total.item()
            rpn_loc_loss += rpn_loc.item()
            rpn_cls_loss += rpn_cls.item()
            roi_loc_loss += roi_loc.item()
            roi_cls_loss += roi_cls.item()

            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()

        acc = model.evaluate(valid_data)
        result = f'loss {total_loss},acc {acc}'
        loss_list.append(total_loss)
        acc_list.append(acc)
        bar.set_description(result)
        printf.info(result)
    return loss_list, acc_list


def test():
    model = config.model()
    model.load()
    acc = model.evaluate(load.test_data)
    printf.debug(f"Accuracy {acc * 100}%")


def train():
    class_names, num_classes = get_classes(config.CLASS_PATH)

    model = FasterRCNN(num_classes)
    # 加载vgg权重
    model_dict = model.state_dict()
    pretrained_dict = torch.load(config.VGG_WEIGHT_PATH, map_location=config.device)
    load_key, no_load_key, temp_dict = [], [], {}
    for k, v in pretrained_dict.items():
        if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
            temp_dict[k] = v
            load_key.append(k)
        else:
            no_load_key.append(k)
    model_dict.update(temp_dict)
    model.load_state_dict(model_dict)

    loss, acc = _train(model, valid_data=load.test_data, train_data=load.train_data)
    save_model(model)
    draw_result(config.epochs, loss, acc)
